# pubmed_searcher.py
# !pip install biopython pyinstaller requests

from Bio import Entrez
import os
import sys
import requests
import time
import random
from datetime import datetime

# 获取应用程序路径，确保在打包后能正确找到资源
def get_application_path():
    """获取应用程序路径，确保在打包后能正确找到资源"""
    if getattr(sys, 'frozen', False):
        # 如果是打包后的exe
        return os.path.dirname(sys.executable)
    else:
        # 如果是脚本运行
        return os.path.dirname(os.path.abspath(__file__))

# 设置你的pubmed邮箱和 API 密钥
Entrez.email = "pengjin@msn.cn"
Entrez.api_key = "1f431**************efd09"

# DeepSeek API配置
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"
DEEPSEEK_MODEL = "deepseek-chat"

def search_pubmed(query_term, max_results=10):
    """
    使用 PubMed API 搜索文献，并返回文献 ID 列表。
    Args:
        query_term: 搜索关键词。
        max_results: 最大返回结果数量。
    Returns:
        文献 ID 列表。
    """
    try:
        handle = Entrez.esearch(db="pubmed", term=query_term, retmax=max_results)
        record = Entrez.read(handle)
        handle.close()
        return record["IdList"]
    except Exception as e:
        print(f"搜索PubMed时出错: {e}")
        return []

def get_paper_details(pmid):
    """
    获取论文的详细信息，包括标题、作者、期刊、发表日期等。
    Args:
        pmid: PubMed ID
    Returns:
        包含论文详细信息的字典
    """
    try:
        handle = Entrez.efetch(db="pubmed", id=pmid, retmode="xml")
        record = Entrez.read(handle)
        handle.close()
        
        article = record['PubmedArticle'][0]['MedlineCitation']['Article']
        journal = article['Journal']
        
        # 获取标题
        title = article['ArticleTitle']
        
        # 获取作者 (如果有)
        authors = []
        if 'AuthorList' in article:
            for author in article['AuthorList']:
                if 'LastName' in author and 'ForeName' in author:
                    authors.append(f"{author['LastName']} {author['ForeName']}")
                elif 'LastName' in author and 'Initials' in author:
                    authors.append(f"{author['LastName']} {author['Initials']}")
                elif 'CollectiveName' in author:
                    authors.append(author['CollectiveName'])
        
        # 获取期刊信息
        journal_name = journal['Title'] if 'Title' in journal else "Unknown Journal"
        
        # 获取发表日期
        pub_date = ""
        if 'PubDate' in journal['JournalIssue']:
            date_parts = journal['JournalIssue']['PubDate']
            if 'Year' in date_parts:
                pub_date = date_parts['Year']
                if 'Month' in date_parts:
                    pub_date = f"{date_parts['Month']} {pub_date}"
                if 'Day' in date_parts:
                    pub_date = f"{date_parts['Day']} {pub_date}"
        
        # 获取DOI (如果有)
        doi = ""
        if 'ELocationID' in article:
            for id in article['ELocationID']:
                if id.attributes['EIdType'] == 'doi':
                    doi = id
        
        # 获取摘要
        abstract = ""
        if 'Abstract' in article and 'AbstractText' in article['Abstract']:
            if isinstance(article['Abstract']['AbstractText'], list):
                for text in article['Abstract']['AbstractText']:
                    if isinstance(text, str):
                        abstract += text + " "
                    else:
                        abstract += text.title() + ": " + text + " "
            else:
                abstract = article['Abstract']['AbstractText']
        
        return {
            "pmid": pmid,
            "title": title,
            "authors": ", ".join(authors),
            "journal": journal_name,
            "pub_date": pub_date,
            "doi": doi,
            "abstract": abstract
        }
    
    except Exception as e:
        print(f"Error getting details for PMID {pmid}: {e}")
        return {
            "pmid": pmid,
            "title": "Error retrieving information",
            "authors": "",
            "journal": "",
            "pub_date": "",
            "doi": "",
            "abstract": ""
        }

def call_deepseek_api(prompt, api_key, max_retries=3):
    """
    调用DeepSeek API生成内容
    Args:
        prompt: 提示词
        api_key: API密钥
        max_retries: 最大重试次数
    Returns:
        生成的文本内容
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": DEEPSEEK_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.7,
        "max_tokens": 2000
    }
    
    for attempt in range(max_retries):
        try:
            response = requests.post(DEEPSEEK_API_URL, headers=headers, json=payload)
            response.raise_for_status()
            data = response.json()
            return data['choices'][0]['message']['content'].strip()
        except Exception as e:
            if attempt < max_retries - 1:
                wait_time = (2 ** attempt) + random.random()
                print(f"API调用失败，等待{wait_time:.2f}秒后重试...")
                time.sleep(wait_time)
            else:
                raise Exception(f"DeepSeek API调用失败: {str(e)}")

def summarize_pubmed_abstract(pmid, deepseek_api_key, query_topic):
    """
    使用 DeepSeek 模型总结 PubMed 文章摘要。
    Args:
        pmid: PubMed ID (例如：'35854100').
        deepseek_api_key: DeepSeek API 密钥
        query_topic: 检索主题，用于提示模型关注相关内容
    Returns:
        总结后的摘要文本，如果发生错误则返回 None。
    """
    try:
        # 从 PubMed 获取文章信息
        paper_details = get_paper_details(pmid)
        abstract = paper_details["abstract"]
        
        if not abstract:
            return {
                "pmid": pmid,
                "title": paper_details["title"],
                "authors": paper_details["authors"],
                "journal": paper_details["journal"],
                "pub_date": paper_details["pub_date"],
                "summary": "未找到摘要。"
            }

        # 使用 DeepSeek 模型进行总结
        prompt = f"""请使用中文总结以下关于"{query_topic}"的PubMed文章摘要。
        包括关键发现、方法和意义，大约3-4个句子:
        
        标题: {paper_details['title']}
        摘要: {abstract}"""
        
        summary = call_deepseek_api(prompt, deepseek_api_key)
        
        return {
            "pmid": pmid,
            "title": paper_details["title"],
            "authors": paper_details["authors"],
            "journal": paper_details["journal"],
            "pub_date": paper_details["pub_date"],
            "summary": summary
        }

    except Exception as e:
        print(f"An error occurred: {e}")
        return {
            "pmid": pmid,
            "title": paper_details.get("title", "Unknown"),
            "authors": paper_details.get("authors", ""),
            "journal": paper_details.get("journal", ""),
            "pub_date": paper_details.get("pub_date", ""),
            "summary": f"无法生成摘要: {str(e)}\n\n原始摘要: {paper_details.get('abstract', '')[:300]}..."
        }

def fetch_abstracts(id_list):
    """
    根据文献 ID 列表，获取文献摘要。
    Args:
        id_list: 文献 ID 列表。
    Returns:
        包含摘要的字符串，每篇摘要之间用换行符分隔。
    """
    if not id_list:
        return ""

    try:
        handle = Entrez.efetch(db="pubmed", id=id_list, idtype="pmid", rettype="abstract", retmode="text")
        abstracts = handle.read()
        handle.close()
        return abstracts
    except Exception as e:
        print(f"获取摘要时出错: {e}")
        return ""

def summarize_abstracts_with_deepseek(abstracts, deepseek_api_key, query_topic):
    """
    使用 DeepSeek 模型，总结多篇文献摘要并生成综合总结。
    Args:
        abstracts: 包含摘要的字符串。
        deepseek_api_key: 您的 DeepSeek API 密钥。
        query_topic: 检索主题
    Returns:
        总结结果字符串。
    """
    try:
        prompt_content = f"""请用中文总结以下关于"{query_topic}"的医学文献摘要。
        分析共同的发现、方法差异、临床意义以及可能的未来研究方向，提炼出最关键的信息:
        
        {abstracts}
        
        请提供一份全面但简明的总结，包括以下方面:
        1. 研究背景和重要性
        2. 主要发现和结论
        3. 临床实践建议
        4. 未来研究方向
        """

        return call_deepseek_api(prompt_content, deepseek_api_key)

    except Exception as e:
        return f"调用 DeepSeek API 出错: {e}。\n\n由于API限制，无法生成综合总结，但您仍可以阅读各文献的原始摘要。"

def process_individual_papers_with_quota_management(pubmed_ids, deepseek_api_key, query_topic, batch_size=3):
    """
    分批处理论文并处理配额限制
    Args:
        pubmed_ids: PubMed ID列表
        deepseek_api_key: DeepSeek API密钥
        query_topic: 检索主题
        batch_size: 每批处理的文献数量
    Returns:
        包含每篇论文摘要的列表
    """
    paper_summaries = []
    
    # 分批处理
    for i in range(0, len(pubmed_ids), batch_size):
        batch_ids = pubmed_ids[i:i+batch_size]
        print(f"\n处理批次 {i//batch_size + 1}/{(len(pubmed_ids) + batch_size - 1)//batch_size}...")
        
        batch_summaries = []
        for j, pmid in enumerate(batch_ids):
            print(f"  处理论文 {i+j+1}/{len(pubmed_ids)}，PMID: {pmid}...")
            
            try:
                summary = summarize_pubmed_abstract(pmid, deepseek_api_key, query_topic)
                batch_summaries.append(summary)
                
                # 打印部分结果
                print(f"  ✓ 成功处理论文 {i+j+1}: {summary['title'][:50]}...\n")
                
                # 添加延迟，避免触发速率限制
                if j < len(batch_ids) - 1:
                    delay = random.uniform(5, 8)  # 随机延迟5-8秒
                    print(f"    等待 {delay:.2f} 秒...")
                    time.sleep(delay)
                    
            except Exception as e:
                error_msg = str(e)
                print(f"  ✗ 处理论文 PMID: {pmid} 时出错: {error_msg}")
                
                # 检查是否是配额限制错误
                if "quota" in error_msg.lower() or "exhausted" in error_msg.lower() or "429" in error_msg:
                    print("\n⚠️ API配额已耗尽，等待较长时间恢复...")
                    
                    # 获取论文基本信息，不使用AI总结
                    try:
                        paper_details = get_paper_details(pmid)
                        batch_summaries.append({
                            "pmid": pmid,
                            "title": paper_details.get("title", "未知标题"),
                            "authors": paper_details.get("authors", ""),
                            "journal": paper_details.get("journal", ""),
                            "pub_date": paper_details.get("pub_date", ""),
                            "summary": "API配额限制：无法生成摘要。原始摘要：\n" + paper_details.get("abstract", "无法获取原始摘要")[:300] + "..."
                        })
                    except:
                        batch_summaries.append({
                            "pmid": pmid,
                            "title": "API配额已耗尽",
                            "authors": "",
                            "journal": "",
                            "pub_date": "",
                            "summary": "由于API配额限制，无法生成摘要。请稍后再试或减少请求数量。"
                        })
                    
                    # 延长等待时间，希望配额会恢复
                    recovery_wait = 60  # 等待60秒
                    print(f"    等待 {recovery_wait} 秒后继续...")
                    time.sleep(recovery_wait)
                else:
                    batch_summaries.append({
                        "pmid": pmid,
                        "title": "处理出错",
                        "authors": "",
                        "journal": "",
                        "pub_date": "",
                        "summary": f"无法生成摘要: {error_msg}"
                    })
        
        # 保存这批处理的结果
        paper_summaries.extend(batch_summaries)
        
        # 批次之间等待较长时间
        if i + batch_size < len(pubmed_ids):
            batch_wait = random.uniform(15, 20)  # 批次间等待15-20秒
            print(f"\n批次处理完成，等待 {batch_wait:.2f} 秒后处理下一批...")
            time.sleep(batch_wait)
    
    return paper_summaries

def create_literature_report(query, pubmed_ids, paper_summaries, overall_summary, output_file=None):
    """
    创建文献报告并保存为文本文件。
    Args:
        query: 搜索查询
        pubmed_ids: PubMed ID列表
        paper_summaries: 每篇论文的摘要
        overall_summary: 总体摘要
        output_file: 输出文件名（可选）
    Returns:
        报告文本内容
    """
    # 如果没有提供输出文件名，生成一个基于日期和查询的文件名
    if output_file is None:
        date_str = datetime.now().strftime("%Y%m%d_%H%M%S")
        sanitized_query = "".join(c if c.isalnum() else "_" for c in query)[:30]
        output_file = os.path.join(get_application_path(), f"PubMed_Report_{sanitized_query}_{date_str}.txt")
    
    # 创建报告内容
    report = []
    
    # 添加标题和搜索信息
    report.append("=" * 80)
    report.append(f"医学文献综述报告: {query}")
    report.append("=" * 80)
    report.append(f"搜索日期: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append(f"搜索关键词: {query}")
    report.append(f"文献数量: {len(pubmed_ids)}")
    report.append("")
    
    # 添加总体摘要
    report.append("=" * 80)
    report.append("综合总结")
    report.append("=" * 80)
    report.append(overall_summary)
    report.append("\n")
    
    # 添加每篇文献的摘要
    report.append("=" * 80)
    report.append("各篇文献摘要")
    report.append("=" * 80)
    
    for i, paper in enumerate(paper_summaries):
        report.append(f"[{i+1}] {paper['title']}")
        report.append(f"作者: {paper['authors']}")
        report.append(f"期刊: {paper['journal']}")
        report.append(f"发表日期: {paper['pub_date']}")
        report.append(f"PubMed ID: {paper['pmid']}")
        report.append("")
        report.append("摘要:")
        report.append(paper['summary'])
        report.append("")
        report.append("-" * 40)
    
    # 添加引用信息
    report.append("=" * 80)
    report.append("引用信息")
    report.append("=" * 80)
    for i, paper in enumerate(paper_summaries):
        report.append(f"[{i+1}] {paper['authors']}. {paper['title']}. {paper['journal']}. {paper['pub_date']}. PMID: {paper['pmid']}")
    
    # 添加生成说明
    report.append("")
    report.append("=" * 80)
    report.append("本报告由 PubMed 文献检索系统自动生成")
    report.append(f"生成日期：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append("=" * 80)
    
    # 将报告内容写入文件
    report_text = "\n".join(report)
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(report_text)
        
        print(f"报告已保存至: {os.path.abspath(output_file)}")
    except Exception as e:
        print(f"保存报告时出错: {e}")
        # 尝试使用备用路径
        try:
            desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")
            backup_file = os.path.join(desktop_path, f"PubMed_Report_{date_str}.txt")
            with open(backup_file, "w", encoding="utf-8") as f:
                f.write(report_text)
            print(f"报告已保存至桌面: {backup_file}")
        except:
            print("无法保存报告。请检查权限或磁盘空间。")
    
    return report_text

def get_user_input():
    """
    获取用户输入的搜索关键词和文献数量
    Returns:
        搜索关键词和文献数量
    """
    print("\n" + "=" * 60)
    print(" " * 15 + "PubMed 文献检索与总结系统")
    print("=" * 60)
    
    # 获取搜索关键词
    while True:
        search_query = input("\n请输入您想检索的医学主题关键词 (建议使用英文): ").strip()
        if search_query:
            break
        print("关键词不能为空，请重新输入！")
    
    # 获取文献数量
    while True:
        try:
            paper_count = input("\n请输入您想检索的文献数量 (1-20): ").strip()
            paper_count = int(paper_count)
            if 1 <= paper_count <= 20:
                break
            elif paper_count > 20:
                print("警告：太花钱了，省着用！请输入不超过20的数字。")
            else:
                print("请输入一个介于1和20之间的整数！")
        except ValueError:
            print("请输入一个有效的整数！")
    
    # 确认信息
    print("\n" + "-" * 60)
    print(f"您将检索关于「{search_query}」的 {paper_count} 篇最新文献")
    print("-" * 60)
    
    return search_query, paper_count

if __name__ == "__main__":
    try:
        # 显示欢迎信息
        print("\n" + "=" * 70)
        print(" " * 20 + "PubMed 文献检索与总结系统")
        print(" " * 15 + "自动检索并总结医学研究文献的最新进展")
        print("=" * 70)
        
        # 询问用户输入API密钥
        your_deepseek_api_key = input("\n请输入您的DeepSeek API密钥 (回车跳过，使用默认密钥): ").strip()
        if not your_deepseek_api_key:
            print("使用默认API密钥。如果遇到错误，请重新运行并提供有效的API密钥。")
            your_deepseek_api_key = "sk-62743778**************66aad"  # 默认密钥
        
        # 获取用户输入
        search_query, paper_count = get_user_input()
        
        # 设置输出文件名（可选，如果不设置，将自动生成）
        output_file = None  # 例如: "医学文献报告.txt"
        
        # 搜索PubMed
        print(f"\n正在从PubMed检索关于「{search_query}」的{paper_count}篇文献...")
        pubmed_ids = search_pubmed(search_query, max_results=paper_count) 

        if pubmed_ids:
            print(f"找到 {len(pubmed_ids)} 篇相关文献，PMID: {pubmed_ids}")
            
            # 显示部分摘要作为示例
            if pubmed_ids:
                first_abstract = fetch_abstracts([pubmed_ids[0]])
                print("\n第一篇文献摘要 (部分):\n", first_abstract[:500] + "...")

            print("\n开始处理各篇论文...")
            # 处理每篇论文 - 使用改进版函数
            paper_summaries = process_individual_papers_with_quota_management(pubmed_ids, your_deepseek_api_key, search_query)
            
            print("\n获取所有论文原始摘要...")
            # 获取所有原始摘要用于综合总结
            all_abstracts = fetch_abstracts(pubmed_ids)
            
            print("\n生成综合总结...")
            # 生成综合总结
            overall_summary = summarize_abstracts_with_deepseek(all_abstracts, your_deepseek_api_key, search_query)
            
            print("\n创建并保存最终报告...")
            # 创建并保存报告
            report = create_literature_report(search_query, pubmed_ids, paper_summaries, overall_summary, output_file)
            
            print("\n报告生成完成！")
            print("\n按任意键退出程序...")
            input()

        else:
            print(f"未找到关于「{search_query}」的相关 PubMed 文献。请尝试使用不同的关键词。")
            print("\n按任意键退出程序...")
            input()
    
    except Exception as e:
        print(f"\n程序运行时出错: {e}")
        print("\n请截图此错误信息并联系开发者。按任意键退出程序...")
        input()